"""
双输出ONNX模型推理示例
用于演示如何使用导出的LCX+LAD双分支分割模型进行推理
"""

import numpy as np
import onnxruntime as ort
import cv2
import os
from typing import Tuple, List
import argparse

class DualOutputONNXPredictor:
    """双输出ONNX模型预测器"""
    
    def __init__(self, onnx_path: str, device: str = 'cpu'):
        """
        初始化预测器
        
        Args:
            onnx_path: ONNX模型文件路径
            device: 推理设备 ('cpu' 或 'cuda')
        """
        self.onnx_path = onnx_path
        
        # 设置推理提供程序
        if device == 'cuda' and 'CUDAExecutionProvider' in ort.get_available_providers():
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
            print("使用GPU进行推理")
        else:
            providers = ['CPUExecutionProvider']
            print("使用CPU进行推理")
        
        # 创建推理会话
        self.session = ort.InferenceSession(onnx_path, providers=providers)
        
        # 获取输入输出信息
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [output.name for output in self.session.get_outputs()]
        
        print(f"输入节点: {self.input_name}")
        print(f"输出节点: {self.output_names}")
        
        # 获取输入尺寸信息
        input_shape = self.session.get_inputs()[0].shape
        self.input_height = int(input_shape[2]) if input_shape[2] != -1 else 512
        self.input_width = int(input_shape[3])   if input_shape[3] != -1 else 512
        print(f"模型输入尺寸: {self.input_height}x{self.input_width}")
    
    def preprocess_image(self, image_path: str) -> np.ndarray:
        """
        预处理输入图像
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            预处理后的图像数组 [1, 3, H, W]
        """
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # BGR转RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 调整尺寸
        image = cv2.resize(image, (int(self.input_width), int(self.input_height)))
        
        # 归一化 (根据训练时的预处理配置)
        # 注意：这里使用与训练配置相同的归一化参数
        image = image.astype(np.float32) / 255.0
        
        # 转换为模型输入格式 [1, 3, H, W]
        image = np.transpose(image, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)   # 添加batch维度
        
        return image
    
    def predict(self, image_path: str) -> Tuple[np.ndarray, np.ndarray]:
        """
        执行推理
        
        Args:
            image_path: 输入图像路径
            
        Returns:
            tuple: (lcx_output, lad_output) 两个分支的输出
        """
        # 预处理图像
        input_data = self.preprocess_image(image_path)
        
        # 执行推理
        outputs = self.session.run(self.output_names, {self.input_name: input_data})
        
        # 通常第一个输出是LCX，第二个是LAD
        lcx_output = outputs[0]  # [1, num_classes, H, W]
        lad_output = outputs[1]  # [1, num_classes, H, W]
        
        return lcx_output, lad_output
    
    def postprocess_outputs(self, lcx_output: np.ndarray, lad_output: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        后处理输出，获取分割掩码
        
        Args:
            lcx_output: LCX分支输出 [1, num_classes, H, W]
            lad_output: LAD分支输出 [1, num_classes, H, W]
            
        Returns:
            tuple: (lcx_mask, lad_mask) 两个分支的分割掩码
        """
        # 移除batch维度并获取预测类别
        lcx_mask = np.argmax(lcx_output[0], axis=0)  # [H, W]
        lad_mask = np.argmax(lad_output[0], axis=0)  # [H, W]
        
        return lcx_mask, lad_mask
    
    def save_results(self, lcx_mask: np.ndarray, lad_mask: np.ndarray, 
                     output_dir: str, filename_prefix: str):
        """
        保存分割结果
        
        Args:
            lcx_mask: LCX分割掩码
            lad_mask: LAD分割掩码
            output_dir: 输出目录
            filename_prefix: 文件名前缀
        """
        os.makedirs(output_dir, exist_ok=True)
        
        # 保存LCX掩码
        lcx_path = os.path.join(output_dir, f"{filename_prefix}_lcx.png")
        cv2.imwrite(lcx_path, (lcx_mask * 255).astype(np.uint8))
        
        # 保存LAD掩码
        lad_path = os.path.join(output_dir, f"{filename_prefix}_lad.png")
        cv2.imwrite(lad_path, (lad_mask * 255).astype(np.uint8))
        
        # 保存合并结果（LCX=1, LAD=2）
        combined_mask = np.zeros_like(lcx_mask)
        combined_mask[lcx_mask > 0] = 1  # LCX区域标记为1
        combined_mask[lad_mask > 0] = 2  # LAD区域标记为2
        
        combined_path = os.path.join(output_dir, f"{filename_prefix}_combined.png")
        cv2.imwrite(combined_path, (combined_mask * 127).astype(np.uint8))  # 0, 127, 254
        
        print(f"结果已保存:")
        print(f"  LCX分割: {lcx_path}")
        print(f"  LAD分割: {lad_path}")
        print(f"  合并结果: {combined_path}")

def main():
    parser = argparse.ArgumentParser(description="双输出ONNX模型推理")
    parser.add_argument('--onnx_path', type=str, default='work_dirs/20250924_170325/onnx/dual_output_model.onnx',
                       help='ONNX模型文件路径')
    parser.add_argument('--image_path', type=str, default='demo.png',
                       help='输入图像路径')
    parser.add_argument('--output_dir', type=str, default='./inference_results',
                       help='输出目录 (默认: ./inference_results)')
    parser.add_argument('--device', type=str, default='cuda', choices=['cpu', 'cuda'],
                       help='推理设备 (默认: cpu)')
    
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.onnx_path):
        print(f"错误: ONNX模型文件不存在: {args.onnx_path}")
        return
    
    if not os.path.exists(args.image_path):
        print(f"错误: 输入图像不存在: {args.image_path}")
        return
    
    try:
        # 创建预测器
        predictor = DualOutputONNXPredictor(args.onnx_path, args.device)
        
        # 执行推理
        print(f"\n开始处理图像: {args.image_path}")
        lcx_output, lad_output = predictor.predict(args.image_path)
        
        print(f"LCX输出形状: {lcx_output.shape}")
        print(f"LAD输出形状: {lad_output.shape}")
        
        # 后处理
        lcx_mask, lad_mask = predictor.postprocess_outputs(lcx_output, lad_output)
        
        # 统计分割结果
        lcx_pixels = np.sum(lcx_mask > 0)
        lad_pixels = np.sum(lad_mask > 0)
        total_pixels = lcx_mask.size
        
        print(f"\n分割统计:")
        print(f"LCX血管像素: {lcx_pixels} ({lcx_pixels/total_pixels*100:.2f}%)")
        print(f"LAD血管像素: {lad_pixels} ({lad_pixels/total_pixels*100:.2f}%)")
        
        # 保存结果
        filename_prefix = os.path.splitext(os.path.basename(args.image_path))[0]
        predictor.save_results(lcx_mask, lad_mask, args.output_dir, filename_prefix)
        
        print(f"\n推理完成! 结果保存在: {args.output_dir}")
        
    except Exception as e:
        print(f"推理失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 